package com.aloudmeta.grpc.client;

import com.aloudmeta.grpc.client.support.ZstdCompressor;
import com.aloudmeta.grpc.client.support.ZstdDecompressor;
import com.aloudmeta.grpc.demo.EchoRequest;
import com.aloudmeta.grpc.demo.EchoResponse;
import com.aloudmeta.grpc.demo.EchoServiceGrpc;
import io.grpc.*;
import io.grpc.stub.MetadataUtils;

import java.util.Iterator;
import java.util.concurrent.TimeUnit;

/**
 * huaixin 2022/1/22 3:54 PM
 */
public class DuuboEchoClient {
    private static Metadata.Key<String> GRPC_TENANT_ID =
            Metadata.Key.of("TENANT_ID", Metadata.ASCII_STRING_MARSHALLER);

    private static Metadata.Key<String> GRPC_USER_ID =
            Metadata.Key.of("USER_ID", Metadata.ASCII_STRING_MARSHALLER);

    public static void main(String[] args) throws InterruptedException {
        CompressorRegistry compressorRegistry = CompressorRegistry.newEmptyInstance();
        compressorRegistry.register(new ZstdCompressor());

        DecompressorRegistry decompressorRegistry = DecompressorRegistry.emptyInstance();
        decompressorRegistry = decompressorRegistry.with(new ZstdDecompressor(), true);

        ManagedChannel originChannel = ManagedChannelBuilder.forAddress("localhost", 50051)
                .usePlaintext()
                .compressorRegistry(compressorRegistry)
                .decompressorRegistry(decompressorRegistry)
                .build();

        ClientInterceptor interceptor = new HeaderClientInterceptor();
        Channel channel = ClientInterceptors.intercept(originChannel, interceptor);

        EchoServiceGrpc.EchoServiceBlockingStub stub = buildStub(channel);
        for (int i = 0; i < 1; i++) {

            try {
                Iterator<EchoResponse> response = stub.echo(EchoRequest.newBuilder()
                        .setMessage("World!")
                        .build());

                while (response.hasNext()) {
                    EchoResponse echoResponse = (EchoResponse) response.next();
                    System.out.println(echoResponse.getMessage());
                }
            } catch (Exception e) {
                e.printStackTrace();
            }
        }

        TimeUnit.SECONDS.sleep(3);
    }

    private static EchoServiceGrpc.EchoServiceBlockingStub buildStub(Channel channel) {
        EchoServiceGrpc.EchoServiceBlockingStub stub = EchoServiceGrpc.newBlockingStub(channel);

        Metadata extraHeaders = new Metadata();
        extraHeaders.put(GRPC_TENANT_ID, "aloudata");
        extraHeaders.put(GRPC_USER_ID, "demo");

        stub = stub.withCompression("zstd");
        stub = stub.withInterceptors(MetadataUtils.newAttachHeadersInterceptor(extraHeaders));
        return stub;
    }
}
